Upgrade to Egglog 3, add extraction controls, and retain Param-Eq demos by saulshanabrook · Pull Request #414 · egraphs-good/egglog-python

saulshanabrook · 2026-04-06T20:19:20Z

Summary

Upgrade the Rust extension and Python API to Egglog 3, expose its newer runtime and extraction controls, and preserve the paused Param-Eq work as a documented experimental module with bounded CI demonstrations.

Runtime and API changes

Breaking: remove Map.rebuild(), Set.rebuild(), and Vec.rebuild(); Egglog 3 rebuilds container values internally.
Breaking: configure worker threads through EGraph(num_threads=...) rather than RAYON_NUM_THREADS.
Breaking: body-defined functions, methods, constants, and class variables lower eagerly unless an explicit ruleset keeps an eqsort body rewrite-backed.
Add Pair, Maybe, catch, map folding, map/set lengths, additional numeric conversions and guarded f64 operations, and an expanded experimental Rational API.
Preserve Python argument order for reverse_args callables and isolate higher-order callable probing from the live ruleset.
Add persistent backoff schedules, RunReport.can_stop, safe naive and opt-in unsafe-seminaive rule evaluation, and per-e-graph/per-rule decomposition controls.
Add source-aware parse-and-run APIs and removable saved transcripts. Replayable failures are recorded with (fail ...); failures that cannot be replayed safely invalidate the transcript.
Propagate Python exceptions raised by PyObject primitives on worker threads back to the caller.
Correct run-report durations that were previously 1,000 times too small, and add var() typing for parameterized expression types.

Maybe, catch, and map folding intentionally remain unsupported in proof mode.

Extraction and costs

Add extractor="tree" | "greedy-dag" to extraction APIs.
Add ordered heterogeneous sequence extraction through extract_multiple, consuming one structured experimental result with a shared TermDag and one variant group per root.
Add destructive keep_best for table-backed constructors, relations, bodyless functions, and bodyless constants.
Use canonical dynamic set_cost tables consistently across extraction, multi-root extraction, and compaction. Negative costs fail when written rather than panicking later.
Rename the general callback protocol to TreeCostModel, retaining CostModel as an alias, and add additive marginal DagCostModel values usable by tree or greedy-DAG extraction.
Breaking: remove GreedyDagCost, GreedyDagCostModel, and greedy_dag_cost_model; use DagCostModel(..., extractor="greedy-dag").

Custom Python cost models remain single-root. A general total tree-cost callback cannot be adapted to greedy-DAG extraction; DagCostModel supplies the required marginal costs.

Param-Eq handoff

egglog.exp.param_eq retains the binary and container representations, restricted expression parser, extraction cost, schedules, CLI, and three project-authored end-to-end examples. Normal pytest/CI runs both representations, parses each extracted result, and compares it with the source at finite sample points. The binary repeated-monomial case remains an explicitly reported iteration-limit stress boundary.

experiments/param_eq contains the optional private-corpus runner, resource guards, provenance validation, and restart notes. Private expressions and external archives are not redistributed. No result CSVs are checked in while the work is paused; any future publication path must generate and validate aggregate-only outputs.

The retained simplifier targets inputs where the source and introduced subexpressions are defined. Its finite sample checks are regression evidence, not a universal equivalence proof; general nonzero/definedness analysis remains future research work.

Dependency stack

egraphs-good/egglog#1008 at 1eea60a14f214505741d22bdd6c9b501b21c04ee
egraphs-good/egglog-experimental#67 at e0a20ce67bbfedf4ace91f4235293165dbb098bc
This PR

Cargo patches the canonical Egglog source so direct bindings and the experimental dependency use the same core revision and Rust type identity.

Validation

uv sync --reinstall-package egglog --all-extras --locked
uv run pytest --benchmark-disable -q - 1,113 passed, 1 skipped, 4 xfailed
make mypy
make stubtest
uv run ruff check .
uv run ruff format --check .
uv lock --check
cargo fmt --check
cargo check --locked --all-targets
cargo test --locked --lib - 4 passed
cargo clippy --locked --all-targets
make docs

The current head is 1e81090581dafce13351b16af404d746dba017d0.

codspeed-hq · 2026-04-06T20:38:30Z

Merging this PR will improve performance by 36.48%

⚠️

Different runtime environments detected

Some benchmarks with significant performance changes were compared across different runtime environments,
which may affect the accuracy of the results.

Open the report in CodSpeed to investigate

⚡ 2 improved benchmarks
✅ 12 untouched benchmarks

Performance Changes

	Mode	Benchmark	`BASE`	`HEAD`	Efficiency
⚡	Simulation	`test_jit[lda]`	7.4 s	5.3 s	+37.83%
⚡	WallTime	`test_jit[lda]`	7.9 s	5.8 s	+35.15%

Tip

Curious why performance improved? Comment @codspeedbot explain why performance improved on this PR, or directly use the CodSpeed MCP with your agent.

_{Comparing sr (1e81090) with main (6b2016e)}

+    def test_init(self):
+        class B(Expr):
+            def __init__(self, value: i64Like) -> None:
+                return B.wrap(value)  # type: ignore[return-value]  # noqa: PLE0101 - symbolic constructor body


Migrate the bindings to Egglog 3, consolidate the general API and correctness fixes with tests and changelog coverage, and preserve Param-Eq as a reusable module and CLI, bounded CI stress cases, and an optional aggregate-only external harness. Pin the Egglog v3 compatibility fix and experimental primitives to immutable revisions, patching the core workspace source so direct and transitive dependencies share one Rust type identity.

saulshanabrook · 2026-08-31T20:57:45Z

+    @method(preserve=True)
+    def pick_key(self) -> T:
+        runtime_self = to_runtime_expr(self)
+        key_type, _value_type = runtime_self.__egg_typed_expr__.tp.args
+        maybe_type = RuntimeClass(
+            Thunk.value(Declarations.create(runtime_self, cast("HasDeclarations", Maybe))),
+            TypeRefWithVars(Ident.builtin("Maybe"), (key_type.to_var(),)),
+            _egg_has_params=True,
+        )
+        initial = cast("Maybe[T]", maybe_type.none())
+        return map_fold_kv(
+            lambda picked, key, _value: picked.match(lambda _: picked, cast("Maybe[T]", maybe_type.some(key))),
+            initial,
+            self,
+        ).unwrap()
+
+    @method(preserve=True)
+    def keys(self) -> Set[T]:
+        runtime_self = to_runtime_expr(self)
+        key_type, _value_type = runtime_self.__egg_typed_expr__.tp.args
+        set_type = RuntimeClass(
+            Thunk.value(Declarations.create(runtime_self, cast("HasDeclarations", Set))),
+            TypeRefWithVars(Ident.builtin("Set"), (key_type.to_var(),)),
+            _egg_has_params=True,
+        )
+        return map_fold_kv(lambda keys, key, _value: keys.insert(key), cast("Set[T]", set_type.empty()), self)


These should be normal expressions, not sure why they are like this... if they aren't builtins then we shouldnt include them here, instead they should just be inlined where we need them... as normal expression, we shouldnt use theprivate runtime expression APIs in public places like this

saulshanabrook · 2026-08-31T20:58:20Z

+def map_filter_kv(f: Callable[[T, V], Unit], xs: Map[T, V]) -> Map[T, V]:
+    runtime_xs = to_runtime_expr(xs)
+    map_type = RuntimeClass(
+        Thunk.value(Declarations.create(runtime_xs, cast("HasDeclarations", Map))),
+        runtime_xs.__egg_typed_expr__.tp.to_var(),
+        _egg_has_params=True,
+    )
+    return map_fold_kv(
+        lambda result, key, value: catch(lambda: f(key, value)).match(lambda _: result.insert(key, value), result),
+        cast("Map[T, V]", map_type.empty()),
+        xs,
+    )
+
+
+def map_map_values(f: Callable[[T, V], V2], xs: Map[T, V]) -> Map[T, V2]:
+    runtime_xs = to_runtime_expr(xs)
+    key_type, value_type = runtime_xs.__egg_typed_expr__.tp.args
+    probe_decls = runtime_xs.__egg_decls__.copy()
+    dummy_key = RuntimeExpr.__from_values__(probe_decls, TypedExprDecl(key_type, DummyDecl()))
+    dummy_value = RuntimeExpr.__from_values__(probe_decls, TypedExprDecl(value_type, DummyDecl()))
+    with set_current_ruleset(None):
+        transformed = cast("Callable[[RuntimeExpr, RuntimeExpr], object]", f)(dummy_key, dummy_value)
+    if not isinstance(transformed, RuntimeExpr):
+        raise TypeError(f"Map value transform must return an egglog expression, got {type(transformed)}")
+    output_type = transformed.__egg_typed_expr__.tp
+    map_type = RuntimeClass(
+        Thunk.value(Declarations.create(runtime_xs, transformed, cast("HasDeclarations", Map))),
+        TypeRefWithVars(Ident.builtin("Map"), (key_type.to_var(), output_type.to_var())),
+        _egg_has_params=True,
+    )
+    return map_fold_kv(
+        lambda result, key, value: result.insert(key, f(key, value)),
+        cast("Map[T, V2]", map_type.empty()),
+        xs,
+    )
+
+
+def map_merge_with(f: Callable[[V, V], V], left: Map[T, V], right: Map[T, V]) -> Map[T, V]:
+    return map_fold_kv(
+        lambda result, key, value: catch(lambda: result[key]).match(
+            lambda old: result.insert(key, f(old, value)), result.insert(key, value)
+        ),
+        left,
+        right,
+    )
+
+


these also shouldnt be here if they are not primitives... they should just be normal expressions where they are used

saulshanabrook · 2026-08-31T21:02:35Z

-    cost_callables: set[CallableRef] = field(default_factory=set)
+    # Cache of top-level expressions lowered with any available synthetic let
+    # references. Rules and rewrites must never read from this cache.
+    expr_to_let_egg_cache: dict[ExprDecl, bindings._Expr] = field(default_factory=dict)


Why is this different from expr_to_egg_cache? Are both required?

saulshanabrook · 2026-08-31T21:04:26Z

+                    # Use constructor declaration instead of constant b/c constants cannot be extracted
+                    # https://github.com/egraphs-good/egglog/issues/334


Constants dont exist anymore

saulshanabrook · 2026-08-31T21:06:55Z

+                from .builtins import ExprValueError  # noqa: PLC0415 - avoid a module import cycle
+                from .runtime import RuntimeExpr  # noqa: PLC0415 - avoid a module import cycle
+
+                if tp.ident == Ident.builtin("Map"):
+                    runtime_expr = RuntimeExpr.__from_values__(self.decls, decl)
+                    try:
+                        as_dict = cast("Map[BaseExpr, BaseExpr]", runtime_expr).value
+                    except ExprValueError:
+                        return self(expr, unwrap_lit=unwrap_lit, ruleset_ident=ruleset_ident, parens=parens), "expr"
+                    if unwrap_lit:
+                        items = ", ".join(
+                            f"{self(cast('RuntimeExpr', k).__egg_typed_expr__, unwrap_lit=True)}: {self(cast('RuntimeExpr', v).__egg_typed_expr__, unwrap_lit=True)}"
+                            for k, v in as_dict.items()
+                        )
+                        return f"{{{items}}}", "Map"
+                    map_str = f"{tp}.empty()"
+                    for key, value in as_dict.items():


We should not use the runtime to rebuild things here... why is this neccessary? Doesnt have to be canoncial, just print expressions as they are without getting value

saulshanabrook · 2026-08-31T21:10:19Z

+    assert '(let $__expr_0 (LetConflictNum_var "explicit"))' in egglog_string
+    assert '(let $__expr_1 (LetConflictNum_var "synthetic"))' in egglog_string


In these tests dont test for the exact string, this seems birttle... also all tests in here when they can should avoid using low level _ APIs and instead just use high level APIs that actually test behavior and not the low level behavior. remove any tests that aren't needed if we just are testing high level behavior

saulshanabrook · 2026-08-31T21:10:52Z

+def test_anonymous_combined_rulesets_use_deterministic_generated_names() -> None:
+    first = ruleset(name="combined_name_probe_first")
+    second = ruleset(name="combined_name_probe_second")
+    combined = unstable_combine_rulesets(first, second)
+    egraph = EGraph(save_egglog_string=True)


like this one too, make sure the tests are verifying a user seeable output, not low level implementaiton details

saulshanabrook · 2026-09-03T16:09:20Z

+## Numeric Predicates and Exact Rationals
+
+The `f64.is_finite()` method returns a `Unit` fact when its value is neither
+infinite nor NaN. This makes it suitable for guarding rules that evaluate
+partial floating-point operations.
+
+The experimental exact `Rational` sort accepts `fractions.Fraction` and
+`i64Like` values in arithmetic, reflected arithmetic, powers, `min`/`max`, and
+ordering predicates. `RationalLike` is the corresponding public type alias.


this seems too specific for this level of doc

saulshanabrook · 2026-09-03T16:09:48Z

-mutate_egraph.register(x)
-mutate_egraph.run(10)
-mutate_egraph.check(eq(x).to(Int(10) + Int(1)))
+incremented = mutate_egraph.let("incremented", x)


why is the run removed here

saulshanabrook · 2026-09-03T16:12:22Z

+Dynamic row costs live in a canonical table named
+`cost_table_<egg-function-name>`. If a compatible, bodyless raw function with
+the same input sorts and `i64` output already has that name when the cost table
+is created, it is reused. An incompatible callable already occupying the name,
+or an incompatible overload that would map to the same canonical table, raises
+an error instead of making the cost table use a generated suffix, because the
+backend only consults the canonical name. Ordinary generated-name collision
+handling still applies to callables registered after the cost table.


this is like mostly low level details, lets not get into this as much

saulshanabrook · 2026-09-03T16:12:50Z

+
+At the low-level bindings layer, the experimental `multi-extract` command
+returns one {class}`egglog.bindings.UserDefinedOutput`.
+{meth}`egglog.bindings.UserDefinedCommandOutput.as_multi_extract` returns a
+{class}`egglog.bindings.MultiExtractOutput` whose `termdag` stores the shared
+term DAG and whose `terms` groups the variant term IDs in root order. It
+returns `None` for a different user-defined output. The high-level method
+performs this conversion automatically.


this is also too low level, we shouldnt talk about bindings here

saulshanabrook · 2026-09-03T16:14:10Z

+Configure worker threads per e-graph with `num_threads`. The default of `1`
+keeps execution serial; `0` uses the machine's available parallelism. You can
+change the setting later with `set_num_threads` and inspect it with
+`num_threads`. The bindings no longer read `RAYON_NUM_THREADS`.


dont reference old env variable

saulshanabrook · 2026-09-03T16:15:02Z

+_CallableMode: TypeAlias = Literal["function", "constructor", "eager", "rewrite"]
+
+
+def _normalize_callable_mode(  # noqa: C901, PLR0911, PLR0912


can you see if this could be reduced in loc somehow? its kinda verbose and hard to follow

saulshanabrook · 2026-09-03T16:16:29Z

+        callback_context = _COST_MODEL_CALLBACK_VALUES.get()
+        in_cost_model_callback = callback_context is not None and callback_context[0] is self
+        if in_cost_model_callback:
+            ref = typed_expr.expr.callable
+            table_is_registered = (
+                ref in self._state.callable_ref_to_egg_fn
+                if isinstance(typed_expr.expr, CallDecl)
+                else ref in self._state.cost_table_names
+            )
+            if not table_is_registered:
+                msg = "Tables queried by cost-model callbacks must be registered before extraction starts"
+                raise ValueError(msg)


this global seems messy why is it needed

saulshanabrook · 2026-09-03T16:17:23Z

+    """
+
+    marginal_cost: Callable[[EGraph, BaseExpr], DAG_COST]
+    identity: DAG_COST


do we really need identity?

+        # already propagating, do not replace it with a cleanup failure.
+        try:
+            call_with_current_trace(self._state.egraph.run_program, bindings.Pop(span(1), 1))
+        except BaseException:


+                "after_params": report.extracted_params,
+            })
+        connection.send(payload)
+    except BaseException:  # worker errors are accounted for without publishing private input text


+            payload = {"status": watched.status}
+        elif process.returncode == 0:
+            payload = _parse_haskell_output(stdout)
+    except (OSError, ValueError):


saulshanabrook force-pushed the sr branch from 5dbbbaa to 9920fd1 Compare April 6, 2026 20:19

github-code-quality Bot found potential problems Apr 6, 2026

View reviewed changes

saulshanabrook force-pushed the sr branch from 009177d to 3b0ef4f Compare April 24, 2026 00:57

github-code-quality Bot found potential problems Apr 24, 2026

View reviewed changes

Comment thread python/egglog/exp/param_eq/pipeline.py Fixed

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

saulshanabrook force-pushed the sr branch from 169f45b to 784b9ab Compare May 18, 2026 20:19

github-code-quality Bot found potential problems May 18, 2026

View reviewed changes

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

saulshanabrook changed the title Sr Upgrade the Python bindings to Egglog 3 and retain Param-Eq stress demos Aug 31, 2026

saulshanabrook force-pushed the sr branch from 22d9066 to ef9dc6e Compare August 31, 2026 20:32

github-code-quality Bot found potential problems Aug 31, 2026

View reviewed changes

saulshanabrook force-pushed the sr branch from ef9dc6e to a822f3f Compare August 31, 2026 20:51

saulshanabrook commented Aug 31, 2026

View reviewed changes

Address Egglog 3 review feedback

2495210

github-code-quality Bot found potential problems Aug 31, 2026

View reviewed changes

Comment thread python/tests/param_eq/test_domain.py Fixed

Comment thread python/egglog/exp/param_eq/domain.py Fixed

Update Egglog dependencies and extraction APIs

6bdaad7

saulshanabrook changed the title ~~Upgrade the Python bindings to Egglog 3 and retain Param-Eq stress demos~~ Upgrade to current Egglog, expose extraction controls, and retain Param-Eq demos Sep 1, 2026

github-code-quality Bot found potential problems Sep 1, 2026

View reviewed changes

Comment thread python/tests/test_high_level.py Fixed

Comment thread python/tests/test_high_level.py Fixed

Comment thread python/tests/test_high_level.py Fixed

Comment thread python/tests/test_high_level.py Fixed

Comment thread python/egglog/egraph.py Fixed

Expose structured multi-extract outputs

ddae949

saulshanabrook commented Sep 3, 2026

View reviewed changes

Finish Egglog 3 integration and Param-Eq handoff

1bb25ea

saulshanabrook changed the title ~~Upgrade to current Egglog, expose extraction controls, and retain Param-Eq demos~~ Upgrade to Egglog 3, add extraction controls, and retain Param-Eq demos Sep 3, 2026

saulshanabrook mentioned this pull request Sep 3, 2026

Make displayed commands replayable and enable output downcasting egraphs-good/egglog#1008

Open

github-code-quality Bot found potential problems Sep 3, 2026

View reviewed changes

Comment thread python/egglog/declarations.py Fixed

Comment thread python/egglog/declarations.py Fixed

Cache explicit backend names during declaration merges

1e81090

github-code-quality Bot found potential problems Sep 3, 2026

View reviewed changes

		# Use constructor declaration instead of constant b/c constants cannot be extracted
		# https://github.com/egraphs-good/egglog/issues/334

		assert '(let $__expr_0 (LetConflictNum_var "explicit"))' in egglog_string
		assert '(let $__expr_1 (LetConflictNum_var "synthetic"))' in egglog_string

		_CallableMode: TypeAlias = Literal["function", "constructor", "eager", "rewrite"]


		def _normalize_callable_mode( # noqa: C901, PLR0911, PLR0912

Uh oh!

Conversation

saulshanabrook commented Apr 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Runtime and API changes

Extraction and costs

Param-Eq handoff

Dependency stack

Validation

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

codspeed-hq Bot commented Apr 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merging this PR will improve performance by 36.48%

Performance Changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

saulshanabrook commented Apr 6, 2026 •

edited

Loading

codspeed-hq Bot commented Apr 6, 2026 •

edited

Loading